#include "shader.h"

#include <fstream>
#include <iostream>
#include <map>
#include <sstream>

#include "renderer.h"

Shader::Shader(const std::string& filepath)
    : filepath_(filepath)
{
    ShaderProgramSource source = ParseShader(filepath);
    rendererID_ = CreateShader(source.vertexSource, source.fragmentSource);
}

Shader::~Shader()
{
    GLCALL(glDeleteProgram(rendererID_));
}

void Shader::Bind() const
{
    GLCALL(glUseProgram(rendererID_));
}

void Shader::Unbind() const
{
    GLCALL(glUseProgram(0));
}

void Shader::SetUniform1i(const std::string& name, int32_t v)
{
    GLCALL(glUniform1i(GetUniformLocation(name), v));
}

void Shader::SetUniform1f(const std::string& name, float v)
{
    GLCALL(glUniform1f(GetUniformLocation(name), v));
}

void Shader::SetUniform4f(const std::string& name, float v0, float v1, float v2, float v3)
{
    GLCALL(glUniform4f(GetUniformLocation(name), v0, v1, v2, v3));
}

void Shader::SetUniformMat4f(const std::string& name, const glm::mat4& matrix)
{
    GLCALL(glUniformMatrix4fv(GetUniformLocation(name), 1, GL_FALSE, &matrix[0][0]));
}

void Shader::SetUniformMat4f(const std::string& name, const Matrix<GLfloat>& matrix)
{
    GLCALL(glUniformMatrix4fv(GetUniformLocation(name), 1, GL_FALSE, matrix));
}

int32_t Shader::GetUniformLocation(const std::string& name)
{
    Bind();
    if (uniformLocationCache_.find(name) != uniformLocationCache_.end()) {
        return uniformLocationCache_[name];
    }
    GLCALL(int32_t location = glGetUniformLocation(rendererID_, name.c_str()));
    if (location == -1) {
        std::cerr << "Warning: uniform '" << name << "' doesn't exist!" << std::endl;
    }

    uniformLocationCache_[name] = location;
    return location;
}

struct ShaderProgramSource Shader::ParseShader(const std::string& filepath)
{
    std::ifstream stream(filepath);
    std::string line;
    std::string type = "";
    std::map<std::string, std::stringstream> ss;

    while (getline(stream, line)) {
        if (line.find("#shader ") == 0) {
            type = line.substr(sizeof("#shader"));
        } else {
            ss[type] << line << "\n";
        }
    }

    return {
        .vertexSource = ss["vertex"].str(),
        .fragmentSource = ss["fragment"].str(),
    };
}

uint32_t Shader::CompileShader(uint32_t type, const std::string& source)
{
    GLCALL(uint32_t id = glCreateShader(type));
    const char *src = source.c_str();
    GLCALL(glShaderSource(id, 1, &src, nullptr));
    GLCALL(glCompileShader(id));

    int32_t result;
    GLCALL(glGetShaderiv(id, GL_COMPILE_STATUS, &result));
    if (result == GL_FALSE) {
        int32_t length;
        GLCALL(glGetShaderiv(id, GL_INFO_LOG_LENGTH, &length));
        char* message = static_cast<char*>(alloca(length));
        GLCALL(glGetShaderInfoLog(id, length, &length, message));
        std::cerr << "CompileShader Failed: " << message;
        GLCALL(glDeleteShader(id));
        return 0;
    }

    return id;
}

uint32_t Shader::CreateShader(const std::string& vertexShader, const std::string& fragmentShader)
{
    GLCALL(uint32_t program = glCreateProgram());
    uint32_t vs = CompileShader(GL_VERTEX_SHADER, vertexShader);
    uint32_t fs = CompileShader(GL_FRAGMENT_SHADER, fragmentShader);

    GLCALL(glAttachShader(program, vs));
    GLCALL(glAttachShader(program, fs));
    GLCALL(glLinkProgram(program));
    GLCALL(glValidateProgram(program));

    GLCALL(glDeleteShader(vs));
    GLCALL(glDeleteShader(fs));
    return program;
}

